from OmniGen import OmniGenPipeline
import random
import numpy as np

pipe = OmniGenPipeline.from_pretrained("/home/bsliu/gitprojects/OmniGen/results/imagenet/checkpoints/0026000", "/home/bsliu/gitprojects/OmniGen/checkpoint/omnigen-v1/vae")  
# pipe.merge_lora("/home/bsliu/gitprojects/OmniGen/results/imagenet_8_cos_32_100/checkpoints/0006000")

# images = pipe(
#     prompt="A player is playing basketball.", 
#     height=1024, 
#     width=1024, 
#     guidance_scale=2.5,
#     seed=0,
# )
# images[0].save("church.png")  # save output PIL Image

hyper_category = [i for i in range(16)]
temp1 = np.exp(-0.2842*np.arange(16))
p = temp1 / np.sum(temp1)

hyper_ids = [("color", [0,1,4,8,9,15]), ("brightness", [2,3]), ("outline", [5,6,7,10,11,12,13,14]), ("color+brightness", [0,1,2,3,4,8,9,15]), ("color+outline",[0,1,4,5,6,7,8,9,10,11,12,13,14,15]), ("brightness+outline", [2,3,5,6,7,10,11,12,13,14]), ("full", [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15])]
single_hc = [(f"{i}", [i]) for i in range(16)]
random_hc = []

for i in range(10):
    bbb = random.sample(hyper_category, random.choices(hyper_category, weights=p, k=1)[0]+1)
    bbb.sort()
    ccc = [str(i) for i in bbb]
    aaa = "-".join(ccc)
    random_hc.append((aaa, bbb))
hyper_ids = hyper_ids + single_hc + random_hc

for item in hyper_ids:
    images = pipe(
        hyper_ids=item[1],
        prompt=f"Perform pattern completion on <img><|image_1|></img>.",
        # input_images=["/share/project/dataset_raw/imagenet/val/n04540053/ILSVRC2012_val_00012002.JPEG"],
        input_images=["/share/project/dataset_raw/imagenet/val/n04591157/ILSVRC2012_val_00009629.JPEG"],
        height=496, 
        width=368,
        guidance_scale=2.5, 
        img_guidance_scale=1.6,
        seed=0
    )
    images[0].save(f"/home/bsliu/gitprojects/OmniGen/imgs/example/hcs_infer2/{item[0]}.png")  # save output PIL image
